package com.mos.iot.mqtt;

import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.StrUtil;
import com.mos.iot.mqtt.callback.MqttClientRoute;
import com.mos.iot.mqtt.config.MqttClientConfig;
import com.mos.iot.mqtt.config.MqttTopicConfig;
import com.mos.iot.mqtt.constants.MqttConstants;
import com.mos.iot.mqtt.listener.MqttTopicMessage;
import lombok.Cleanup;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.eclipse.paho.client.mqttv3.*;
import org.eclipse.paho.client.mqttv3.persist.MemoryPersistence;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Mono;
import reactor.core.publisher.MonoSink;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
@Component
@RequiredArgsConstructor
public class MqttClientManager {

    /**
     * 存储MQTT客户端
     */
    public static Map<String, MqttClient> MQTT_CLIENT_MAP = new ConcurrentHashMap<>();

    public static Map<String, AtomicInteger> MQTT_CLIENT_SERIALNO_MAP = new ConcurrentHashMap<>();

    public static final Map<String, MonoSink<Object>> topicResponse = new HashMap<>();

    public static void addMqttClient(MqttClientConfig mqttClientConfig) {
        addMqttClient(mqttClientConfig.getBroker(),
                mqttClientConfig.getClientId(),
                mqttClientConfig.getUserName(),
                mqttClientConfig.getPassword(),
                mqttClientConfig.getTopics());
    }

    /**
     * 添加MQTT客户端到某个集合中，用于后续的连接和订阅操作。
     *
     * @param broker MQTT代理服务器的地址，用于客户端连接。
     * @param clientId 客户端的唯一标识，同一个客户端在一次会话中应该使用相同的clientId。
     * @param userName MQTT服务器登录时的用户名，如果不需要认证则可以为空。
     * @param password MQTT服务器登录时的密码，如果不需要认证则可以为空。
     * @param topics MQTT订阅主题
     */
    public static void addMqttClient(String broker, String clientId, String userName, String password, List<MqttTopicConfig> topics) {
        try {
            @Cleanup
            MemoryPersistence persistence = new MemoryPersistence();
            MqttClient client = new MqttClient(broker, clientId, persistence);
            MqttConnectOptions connOpts = new MqttConnectOptions();
            connOpts.setCleanSession(true);
            if(StringUtils.hasText(userName)) {
                connOpts.setUserName(userName);
            }
            if(StringUtils.hasText(password)) {
                connOpts.setPassword(password.toCharArray());
            }
            client.setCallback(new MqttClientRoute(clientId, connOpts));
            client.connect(connOpts);
            log.info("mqtt客户端[{}]连接完成。", clientId);
            if(topics != null) {
                for (MqttTopicConfig topicConfig : topics) {
                    client.subscribe(topicConfig.getTopic(), topicConfig.getQos());
                    log.info("mqtt客户端[{}] --> 订阅主题[{}],QOS={};", clientId, topicConfig.getTopic(), topicConfig.getQos());
                }
            }
            MQTT_CLIENT_MAP.putIfAbsent(clientId, client);
            MQTT_CLIENT_SERIALNO_MAP.putIfAbsent(clientId, new AtomicInteger(1));
        } catch (Exception e) {
            log.error("Create mqttClient failed!", e);
        }
    }

    public static MqttClient getMqttClient(String clientId) {
        return MQTT_CLIENT_MAP.get(clientId);
    }

    public static Integer getSerialNo(String clientId) {
        MQTT_CLIENT_SERIALNO_MAP.putIfAbsent(clientId, new AtomicInteger(1));
        AtomicInteger atomicInteger = MQTT_CLIENT_SERIALNO_MAP.get(clientId);
        if(atomicInteger.get() == 65535) {
            atomicInteger.set(1);
        }
        return atomicInteger.getAndIncrement();
    }

    public static <T> T syncPublish(MqttTopicMessage message, String responseKey, int timeout) {
        if(message != null) {
            Mono<Object> receive = subscribe(responseKey);
            if (receive != null) {
                try {
                    return (T)Mono.create(sink -> {
                                publish(message);
                                sink.success();
                            })
                            .then(receive)
                            .timeout(Duration.ofSeconds(timeout))
                            .doFinally(signal -> unsubscribe(responseKey))
                            .block();
                } catch (Exception e) {
                    if(e.getCause() instanceof TimeoutException) {
                        throw new RuntimeException("超时响应", e.getCause());
                    }
                    log.error("syncPublish error", e.getCause());
                }
            }
        }
        return null;
    }

    public void response(String responseKey, Object object) {
        MonoSink<Object> sink = topicResponse.get(responseKey);
        if (sink != null) {
            sink.success(object);
        }
    }

    public static void publish(MqttTopicMessage message) {
        if(message != null) {
            publish(message.getClientId(), message.getTopic(), message.getQos(), message.isRetained(), message.getOutMessage());
        }
    }

    public static void publish(String clientId, String topic, int qos, boolean retained, MqttMessage message) {
        MqttClient mqttClient = getMqttClient(StrUtil.blankToDefault(clientId, MqttConstants.DEFAULT_CLIENT_ID));
        Assert.notNull(mqttClient, "[{}]通道不存在!", clientId);
        Assert.state(mqttClient.isConnected(), "[{}]通道不可用!");
        Assert.notBlank(topic, "主题不能为空!");
        Assert.notNull(message, "消息不能为空!");
        message.setQos(qos);
        message.setRetained(retained);
        MqttTopic mqttTopic = mqttClient.getTopic(topic);
        Assert.notNull(mqttTopic, "[{}]主题未订阅!", topic);
        try {
            mqttTopic.publish(message).waitForCompletion();
            log.info("mqtt客户端[{}] --> 主题[{}]消息发送成功 --> {}", clientId, topic, new String(message.getPayload()));
        } catch (Exception e) {
            throw new RuntimeException("mqtt消息推送失败!", e);
        }
    }

    private static Mono<Object> subscribe(String key) {
        synchronized (topicResponse) {
            if (!topicResponse.containsKey(key)) return Mono.create(sink -> topicResponse.put(key, sink));
        }
        return null;
    }

    private static void unsubscribe(String key) {
        topicResponse.remove(key);
    }

}
